import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
from torch import optim
from torch.utils.data import DataLoader, Dataset


def new_LSTM(input_size, hidden_size, num_layers, output_size):
    model = LSTMMODEL(input_size, hidden_size, num_layers, output_size=output_size)
    return model


class LSTMMODEL(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size=1):
        super(LSTMMODEL, self).__init__()  # 继承父类的属性
        # 定义LSTM的第一层：batch_first=True表示输入数据格式为(batch_size, seq_len, features) 【叶行列】
        self.lstm1 = nn.LSTM(input_size, hidden_size, batch_first=True, num_layers=num_layers, dropout=0.2)
        # 定义LSTM的第二层.
        self.lstm2 = nn.LSTM(hidden_size, hidden_size, batch_first=True, num_layers=num_layers, dropout=0.2)
        # 定义全连接层1，将隐藏层输出转换为10个类别
        self.fc1 = nn.Linear(hidden_size, 10)
        # 定义全连接层2，将10维映射到输出层1维
        self.fc2 = nn.Linear(10, output_size)

    def forward(self, x):
        x, _ = self.lstm1(x)
        x, _ = self.lstm2(x)  # lstm2处理lstm1的输出
        x = self.fc1(x[:, -1, :])
        x = self.fc2(x)
        return x
